// src/server.rs
use tokio::net::{TcpListener, TcpStream};
use tokio_tungstenite::accept_async;
use serde_json::Value;
use rustls::ServerConfig;
use std::fs::File;
use std::io::BufReader;
use std::sync::Arc;
use tokio_rustls::TlsAcceptor;
use crate::transport::{Transport, WebSocketTransport, TcpTransport};
use crate::client::protocol;
use async_trait::async_trait;

struct BufferedTransport {
    inner: Box<dyn Transport + Send>,
    buffer: Vec<u8>,
}

impl BufferedTransport {
    async fn read_exact(&mut self, len: usize) -> anyhow::Result<Vec<u8>> {
        while self.buffer.len() < len {
            let chunk = self.inner.read().await?.ok_or_else(|| anyhow::anyhow!("early eof"))?;
            self.buffer.extend_from_slice(&chunk);
        }
        Ok(self.buffer.drain(..len).collect())
    }
}

#[async_trait]
impl Transport for BufferedTransport {
    async fn read(&mut self) -> anyhow::Result<Option<Vec<u8>>> {
        if !self.buffer.is_empty() {
            let data = self.buffer.clone();
            self.buffer.clear();
            return Ok(Some(data));
        }
        self.inner.read().await
    }

    async fn write(&mut self, data: &[u8]) -> anyhow::Result<()> {
        self.inner.write(data).await
    }

    fn is_closed(&self) -> bool {
        self.inner.is_closed()
    }

    async fn shutdown(&mut self) -> anyhow::Result<()> {
        self.inner.shutdown().await
    }
}

trait IoStream: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin + Send {}
impl<T: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin + Send> IoStream for T {}

#[derive(Debug)]
pub struct Config {
    pub listen_addr: String,
    pub cert: String,
    pub key: String,
    pub transport: String,
    pub no_tls: bool,
}

pub async fn start_server(config: Config, mut shutdown_receiver: tokio::sync::oneshot::Receiver<()>, addr_sender: tokio::sync::oneshot::Sender<std::net::SocketAddr>) -> anyhow::Result<()> {
    let listener = TcpListener::bind(&config.listen_addr).await?;
    let addr = listener.local_addr()?;
    let _ = addr_sender.send(addr);
    log::info!("Server listening on {} with {} transport ", addr, config.transport);

    let tls_acceptor = if !config.no_tls && config.transport == "ws" {
        let certs = rustls_pemfile::certs(&mut BufReader::new(File::open(&config.cert)?)).collect::<Result<Vec<_>, _>>()?;
        let key = rustls_pemfile::private_key(&mut BufReader::new(File::open(&config.key)?))?.ok_or_else(|| anyhow::anyhow!("No private key found"))?;
        let server_config = rustls::ServerConfig::builder_with_provider(
            rustls::crypto::ring::default_provider().into()
        )
        .with_protocol_versions(&[&rustls::version::TLS13])
        .unwrap()
        .with_no_client_auth()
        .with_single_cert(certs, key)?;
        Some(TlsAcceptor::from(Arc::new(server_config)))
    } else {
        None
    };

    loop {
        tokio::select! {
            _ = &mut shutdown_receiver => {
                log::info!("Server shutting down.");
                break;
            }
            conn = listener.accept() => {
                let (stream, _) = conn?;
                let config_transport = config.transport.clone();
                let tls_acceptor = tls_acceptor.clone();

                tokio::spawn(async move {
                    let result = async {
                        let mut transport: Box<dyn Transport + Send> = {
                            let config_transport = config_transport.clone();
                            let tls_acceptor = tls_acceptor.clone();
                            if config_transport == "ws" {
                                let stream: Box<dyn IoStream> = if let Some(acceptor) = tls_acceptor {
                                    Box::new(acceptor.accept(stream).await?)
                                } else {
                                    Box::new(stream)
                                };
                                let ws_stream = accept_async(stream).await?;
                                log::info!("Accepted WebSocket connection");
                                Box::new(WebSocketTransport::new(ws_stream))
                            } else {
                                log::info!("Accepted TCP connection");
                                Box::new(TcpTransport::new(stream))
                            }
                        };

                        let (target_host, target_port) = if config_transport == "ws" {
                            if let Some(data) = transport.read().await? {
                                let json: Value = serde_json::from_slice(&data)?;
                                let host = json.get("host").and_then(Value::as_str).ok_or_else(|| anyhow::anyhow!("Missing or invalid host in metadata"))?.to_string();
                                let port = json.get("port").and_then(Value::as_u64).ok_or_else(|| anyhow::anyhow!("Missing or invalid port in metadata"))? as u16;
                                (host, port)
                            } else {
                                return Err(anyhow::anyhow!("Failed to read metadata from transport"));
                            }
                        } else {
                            let mut buffered_transport = BufferedTransport {
                                inner: transport,
                                buffer: Vec::new(),
                            };
                            let len_bytes = buffered_transport.read_exact(4).await?;
                            let len = u32::from_be_bytes(len_bytes.try_into().unwrap()) as usize;
                            let metadata_bytes = buffered_transport.read_exact(len).await?;
                            let json: Value = serde_json::from_slice(&metadata_bytes)?;
                            let host = json.get("host").and_then(Value::as_str).ok_or_else(|| anyhow::anyhow!("Missing or invalid host in metadata"))?.to_string();
                            let port = json.get("port").and_then(Value::as_u64).ok_or_else(|| anyhow::anyhow!("Missing or invalid port in metadata"))? as u16;
                            transport = Box::new(buffered_transport);
                            (host, port)
                        };

                        let target_stream = TcpStream::connect(format!("{}:{}", target_host, target_port)).await?;
                        log::info!("Connected to target: {}:{}", target_host, target_port);

                        protocol::forward_streams(target_stream, transport).await
                    };

                    if let Err(e) = result.await {
                        log::error!("Error processing server connection: {}", e);
                    }
                });
            }
        }
    }
    Ok(())
}


#[cfg(test)]
mod tests {
    use super::*;
    use crate::client;
    use tokio::io::{AsyncReadExt, AsyncWriteExt};

    // 1. Target Echo Server
    async fn spawn_target_echo_server() -> anyhow::Result<std::net::SocketAddr> {
        let listener = TcpListener::bind("127.0.0.1:0").await?;
        let addr = listener.local_addr()?;
        tokio::spawn(async move {
            let (mut stream, _) = listener.accept().await.unwrap();
            let (mut reader, mut writer) = stream.split();
            tokio::io::copy(&mut reader, &mut writer).await.unwrap();
        });
        Ok(addr)
    }

    // 2. Test runner
    async fn run_e2e_test(transport: &str, no_tls: bool) {
        // Step 1: Start Target Server
        let target_addr = spawn_target_echo_server().await.unwrap();

        // Step 2: Start sshWs Server
        let server_config = Config {
            listen_addr: "127.0.0.1:0".to_string(),
            cert: "cert.pem".to_string(),
            key: "key.pem".to_string(),
            transport: transport.to_string(),
            no_tls,
        };
        let (server_shutdown_sender, server_shutdown_receiver) = tokio::sync::oneshot::channel();
        let (server_addr_sender, server_addr_receiver) = tokio::sync::oneshot::channel();
        tokio::spawn(start_server(server_config, server_shutdown_receiver, server_addr_sender));
        let server_addr = server_addr_receiver.await.unwrap();
        
        // 3. Start sshWs Client
        let client_config = client::Config {
            listen_addr: "127.0.0.1:0".to_string(),
            server_addr: server_addr.to_string(),
            transport: transport.to_string(),
            no_tls,
        };
        let (client_shutdown_sender, client_shutdown_receiver) = tokio::sync::oneshot::channel();
        let (client_addr_sender, client_addr_receiver) = tokio::sync::oneshot::channel();
        tokio::spawn(client::start_client(client_config, client_shutdown_receiver, client_addr_sender));
        let client_addr = client_addr_receiver.await.unwrap();

        // Wait for the client server to start
        let mut attempts = 0;
        while let Err(_) = TcpStream::connect(client_addr).await {
            attempts += 1;
            if attempts > 100 {
                panic!("Client server did not start in time");
            }
            tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;
        }

        // 4. Simulate user application connecting to the SOCKS5 proxy
        let mut user_stream = TcpStream::connect(client_addr).await.unwrap();

        // SOCKS5 Handshake
        user_stream.write_all(&[0x05, 0x01, 0x00]).await.unwrap();
        let mut resp = [0u8; 2];
        user_stream.read_exact(&mut resp).await.unwrap();
        assert_eq!(resp, [0x05, 0x00]);

        // SOCKS5 Request to connect to the target server
        let mut req = vec![0x05, 0x01, 0x00]; // SOCKS5, CONNECT, RSV
        match target_addr.ip() {
            std::net::IpAddr::V4(ip) => {
                req.push(0x01); // ATYP: IPv4
                req.extend(ip.octets());
            }
            std::net::IpAddr::V6(ip) => {
                req.push(0x04); // ATYP: IPv6
                req.extend(ip.octets());
            }
        }
        req.extend(target_addr.port().to_be_bytes());
        user_stream.write_all(&req).await.unwrap();
        let mut resp = [0u8; 10];
        user_stream.read_exact(&mut resp).await.unwrap();
        assert_eq!(resp[1], 0x00, "SOCKS5 request failed");

        // 5. Send data and assert echo
        user_stream.write_all(b"end-to-end").await.unwrap();
        let mut buf = [0u8; 10];
        user_stream.read_exact(&mut buf).await.unwrap();
        assert_eq!(&buf, b"end-to-end");

        // Shut down servers
        server_shutdown_sender.send(()).unwrap();
        client_shutdown_sender.send(()).unwrap();
    }

    #[tokio::test]
    async fn test_e2e_tcp() {
        run_e2e_test("tcp", true).await;
    }

    #[tokio::test]
    async fn test_e2e_websocket_no_tls() {
        run_e2e_test("ws", true).await;
    }

    #[tokio::test]
    async fn test_e2e_websocket_tls() {
        run_e2e_test("ws", false).await;
    }
}
